
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import classification_report
from sklearn.tree import export_graphviz
import graphviz
import pandas as pd

# 读取数据
data = pd.read_csv(r'titanic_data.csv')
data.drop('PassengerId', axis=1, inplace=True)
data.loc[data['Sex'] == 'male', 'Sex'] = 1
data.loc[data['Sex'] == 'female', 'Sex'] = 0
data.fillna(data['Age'].mean(), inplace=True)

Dtc = DecisionTreeClassifier(max_depth=5, random_state=8)
Dtc.fit(data.iloc[:, 1:], data['Survived'])
pre = Dtc.predict(data.iloc[:, 1:])

print(classification_report(data['Survived'], pre))

dot_data = export_graphviz(
    Dtc,
    feature_names=['Pclass', 'Sex', 'Age'],
    class_names='Survived'
)

graph = graphviz.Source(dot_data)
graph.view()


